import datetime
import functools
import logging
import os
import traceback
from concurrent.futures import ProcessPoolExecutor

from absl import app, flags
import pandas as pd

import cc.appcoin2.strategy.run_sim as run_sim
import coin.research.read_strat_info as rsi
from coin.base.datetime_util import (
    to_timestamp_int,
    convert_string_to_date,
    iterate_date)
from coin.base.param_util import to_list
from coin.base.query_util import query_pta
from coin.proto.coin_request_pb2 import (
    AccountRequestProto,
    ProcessInfoProto,
    StrategyRequestProto)
from fastfeature.simulator_spec_pb2 import OrderResponsesProto
from coin.support.pta.app.sim_result_converter import convert_sim_result_to_strat_log


def _convert_sim_type(sim_type):
  converted = None
  if sim_type == 'lm_sim':
    converted = 'lm'
  elif sim_type.startswith('basis_smm'):
    converted = 'basis_smm'
  else:
    raise ValueError('Unknown sim type', sim_type)
  return converted


def get_daily_strategy_info(start_date, end_date):
  fetch_cols = ['trading_date', 'strategy_name', 'machine',
                'market_type', 'exchange', 'owner']
  daily_strategy_info = {}
  pta_df = query_pta(
      start_date, end_date, business_units=['Coin'], pta_type='SYMBOL')[fetch_cols]
  for trading_date, group in pta_df.groupby('trading_date'):
    daily_strategy_info[trading_date] = group.drop_duplicates(
        ['strategy_name', 'machine'])
  return daily_strategy_info


def launch_sim(strat_conf):
  sim_responses = []
  sparams = [{}]
  num_cpu = 20
  stratfile, sim_type = strat_conf['stratfile'], strat_conf['sim_type']
  sim_type = _convert_sim_type(sim_type)
  for sim_backtests_each in run_sim.run_batch_for_params(
      sim_type, to_list(stratfile), sparams, run_sim.batch_local, num_cpu):
    sim_responses.extend(sim_backtests_each)
  return sim_responses


def main(_):
  FLAGS = flags.FLAGS
  output_dir = os.path.abspath(
      os.path.normpath(os.path.expanduser(FLAGS.strat_log_dir)))
  cmd_relpath = os.path.relpath(os.path.abspath(__file__))
  if cmd_relpath.startswith("python/"):
    root_path = "python"
  else:
    root_path = "."
  start_date = convert_string_to_date(FLAGS.trading_date)
  end_date = start_date + datetime.timedelta(days=1)
  trading_dates = [td for td in iterate_date(start_date, end_date)]
  daily_strategy_info = get_daily_strategy_info(
      start_date, end_date - datetime.timedelta(days=1))
  name_conf_map = rsi.get_strat_name_conf_map(root_path)

  for trading_date in trading_dates:
    td_str = trading_date.strftime('%Y%m%d')
    FLAGS.trading_date = td_str
    logging.info('Running for %s ...' % td_str)
    logging.info('Start running sim ...')
    sim_responses = []
    future_list = []
    sim_fail_list = []
    conversion_fail_list = []
    with ProcessPoolExecutor(max_workers=FLAGS.max_workers) as executor:
      for _, strategy_info in daily_strategy_info.get(
          trading_date, pd.DataFrame()).iterrows():
        strategy_name = strategy_info['strategy_name']
        # Now only coin2 sim is supported
        if not strategy_name.startswith(
            tuple(['basis_smm2', 'lm_agg2', 'lm_pass2', 'interval'])):
          continue
        if strategy_name not in name_conf_map:
          logging.info("omitting %s" % strategy_name)
          continue
        future = executor.submit(
            functools.partial(launch_sim, name_conf_map[strategy_name]))
        future_list.append((strategy_info, future))
    for strategy_info, future in future_list:
      key = (strategy_info['machine'], strategy_info['strategy_name'])
      try:
        sim_responses.append((strategy_info, future.result()))
      except Exception as e:
        sim_fail_list.append(key)
        logging.info('Fail to run sim. %s, %s, %s', key, type(e), e)
        logging.info(traceback.format_exc())

    logging.info('Start running strat log conversion ...')
    for strategy_info, sim_backtests in sim_responses:
      key = (strategy_info['machine'], strategy_info['strategy_name'])
      try:
        assert len(sim_backtests) == 1
        sim_backtest = sim_backtests[0]
        strategy_request = StrategyRequestProto(
            strategy_name=strategy_info['strategy_name'])
        process_info = ProcessInfoProto(machine=strategy_info['machine'])
        account_request = AccountRequestProto(
            market_type=strategy_info['market_type'],
            exchange=strategy_info['exchange'],
            owner=strategy_info['owner'])
        sim_log_dir = sim_backtest.tmp_dirname
        pb_files = [filename for filename in os.listdir(sim_log_dir)
                    if filename.endswith('.pb')]
        sim_results = []
        for pb_file in pb_files:
          with open(os.path.join(sim_log_dir, pb_file), 'rb') as fread:
            sim_result = OrderResponsesProto()
            sim_result.ParseFromString(fread.read())
            sim_results.append(sim_result)
        convert_sim_result_to_strat_log(
            trading_date=trading_date, strategy_request=strategy_request,
            process_info=process_info, account_request=account_request,
            sim_results=sim_results, output_dir=output_dir)
      except Exception as e:
        conversion_fail_list.append(key)
        logging.info('Fail to convert strat log. %s, %s, %s', key, type(e), e)
        logging.info(traceback.format_exc())

    if len(sim_fail_list) > 0:
      logging.info('sim_fail_strats:\n%s', sim_fail_list)
    if len(conversion_fail_list) > 0:
      logging.info('conversion_fail_strats:\n%s', conversion_fail_list)
    return len(sim_fail_list) + len(conversion_fail_list)


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', None, '')
  flags.DEFINE_string('strat_prefix', None, '')
  flags.DEFINE_string('pick_symbol_regex', None, '')
  flags.DEFINE_string('analyzer_feed_machine', '', '')
  flags.DEFINE_string('hours', None, '')
  flags.DEFINE_string('sim_result_dir', 'sim_test2', '')
  flags.DEFINE_bool('recompile', None, '')
  flags.DEFINE_integer('group_dates', None, 'if 3, 3 days sim per once')
  flags.DEFINE_string('split_hours', None, '')
  flags.DEFINE_bool('debug', False, '')
  flags.DEFINE_bool('use_live_latency', False, '')
  flags.DEFINE_bool('use_latency_model', False, '')
  flags.DEFINE_bool('exclude_trading_halt', False,
                    'exclude trading halt time in interval')
  flags.DEFINE_string('exclude_date', None, '')
  flags.DEFINE_integer('max_workers', 24, 'Max number of workers.')
  flags.DEFINE_string('strat_log_dir', '~/data/sim_strat_proto_log', '')


  logging.basicConfig(
      level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
